import asyncio
import os
import dotenv

from openai import AsyncOpenAI

from agents import (
    Agent,
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
    RunConfig,
    Runner,
    function_tool,
    set_tracing_disabled,
)

dotenv.load_dotenv()

client = AsyncOpenAI(api_key=os.getenv("AI_API_KEY"),base_url=os.getenv("AI_BASE_URL"))
set_tracing_disabled(disabled=True)

class CustomModelProvider(ModelProvider):
    def get_model(self, model_name: str | None) -> Model:
        return OpenAIChatCompletionsModel(model=os.getenv("AI_API_MODEL"), openai_client=client)

CUSTOM_MODEL_PROVIDER = CustomModelProvider()

# 定义多个agent
历史辅导员  = Agent(
    name="历史辅导员",
    instructions="""
    你是一个历史辅导员，你的任务是帮助用户解决历史事件、历史人物、历史朝代等相关问题。
    你只能回答与历史事件、历史人物、历史朝代直接相关的问题。
    如果用户的问题不是历史事件、历史人物、历史朝代相关的问题，请直接回答：臣妾做不到啊。
    每一步都要解释你的推理过程并包含例子。
    """,
    handoff_description="""
    专门处理历史事件、历史人物、历史朝代等相关问题的 agent
    """
)
数学辅导员 = Agent(
    name="数学辅导员",
    handoff_description="专门处理数学问题的 agent",
    instructions="你提供数学问题的帮助。每一步都要解释你的推理过程并包含例子",
)

#分类处理 Agent
triage_agent = Agent(
    name="分类 Agent",
    instructions="""
    你根据用户的作业问题决定使用哪个 agent。你只能在[历史辅导员, 数学辅导员]中选择，不能创造新的 agent。
    你必须先输出 transfer_to_XXX（XXX为具体agent名称），然后再输出答案。
    例如：
    - 用户问“清朝有几个皇帝”，你应该先输出 transfer_to_历史辅导员，然后再输出具体答案。
    - 用户问“1+1等于几”，你应该先输出 transfer_to_数学辅导员，然后再输出具体答案。
    - 用户问“股票的真谛是什么”、“你好翻译成英文”、“命里有时终需有，下一句是”，这些都不属于历史或数学，你应该只输出‘臣妾做不到啊’。
    如果没有找到对应的agent，请只输出‘臣妾做不到啊’。
    """,
    handoffs=[历史辅导员, 数学辅导员]
)
async def main():
    user_input = input("请你要问的问题: ")
    result = await Runner.run(
        starting_agent = triage_agent,
        input =user_input,
        run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),
    )
    print(result.final_output)

if __name__ == "__main__":  
    asyncio.run(main())


